import os
import uuid

from dotenv import load_dotenv
from flask import Flask, request, jsonify, send_file
from cozepy import Coze, TokenAuth, COZE_CN_BASE_URL, Message, ChatStatus

# 加载环境变量
load_dotenv()

# 创建flask应用实例
app = Flask(__name__)

# 用于存储会话的字典，key为客户端或自定义标识，value为会话信息
user_sessions = {}

# 定义COZE服务
class CozeService:
    def __init__(self):
        # 访问COZE API的令牌
        self.api_token = os.getenv('COZE_API_TOKEN')
        # 要访问的智能体的ID
        self.bot_id = os.getenv('BOT_ID')
        # 访问用户的ID
        self.user_id = os.getenv('USER_ID')

        # 初始化COZE客户端
        self.coze = Coze(
            auth=TokenAuth(token=self.api_token),
            base_url=COZE_CN_BASE_URL
        )

    # 获取智能体的响应
    def get_sdk_resp(self, user_message, user_session):

        try:
            # 判断是新会话，还是已有会话

            # 已有会话记录
            if user_session in user_sessions:
                session_data = user_sessions[user_session]
                conversation_id = session_data['conversation_id']
                user_id = session_data['user_id']

                # 构建用户要发送的消息
                messages = [
                    Message(
                        role="user",
                        content=user_message,
                        content_type="text",
                        type="question"
                    )
                ]

                # 与智能体进行对话
                chat = self.coze.chat.create(
                    bot_id=self.bot_id,
                    user_id=user_id,
                    conversation_id=conversation_id,
                    additional_messages=messages,
                    auto_save_history=True
                )

            else:
                user_id = str(uuid.uuid4())
                # 构建用户要发送的消息
                messages = [
                    Message(
                        role="user",
                        content=user_message,
                        content_type="text",
                        type="question"
                    )
                ]

                # 创建新的会话，与智能体进行对话
                chat = self.coze.chat.create(
                    bot_id=self.bot_id,
                    user_id=user_id,
                    additional_messages=messages,
                    auto_save_history=True
                )

                # 把新加的信息都记录下来
                user_sessions[user_session] = {
                    "conversation_id": chat.conversation_id,
                    "user_id": user_id,
                    "chat_id": chat.id
                }


            # 等待对话结果
            while chat.status == ChatStatus.IN_PROGRESS:
                chat = self.coze.chat.retrieve(
                    conversation_id=chat.conversation_id,
                    chat_id=chat.id
                )
            # 获取智能体的消息
            if chat.status == ChatStatus.COMPLETED:
                messages=self.coze.chat.messages.list(
                    conversation_id=chat.conversation_id,
                    chat_id=chat.id
                )

                for msg in messages:
                    if msg.role == 'assistant':
                        return {"status":"success", "content":msg.content}

            return {"status":"success", "content":"对话失败"}

        except Exception as e:
            return {"status":"error", "content": str(e)}


coze_service = CozeService()

@app.route('/')
def view():
    return send_file('index.html')

# 申明接口路由 post + json
@app.route('/chat', methods=['POST'])
def chat():
    # 获取请求参数
    data = request.json
    user_message = data.get('message')

    user_session = request.remote_addr

    result = coze_service.get_sdk_resp(user_message, user_session)
    return jsonify(result)

# 清理会话的函数
def clean():
    expired_sessions = []
    for user_session, session_data in user_sessions.items():
        expired_sessions.append(user_session)

    for user_session in expired_sessions:
        del user_sessions[user_session]

# 启动服务
if __name__ == '__main__':
    app.run(debug=True, port=5000)
